package com.will.yuaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.will.yuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * @author will
 * @since 2025/7/24 19:01
 * 抽象的代理类,用于管理代理和执行流程
 */
@Data
@Slf4j
public abstract class BaseAgent {

    //核心属性
    private String name;

    //提示词
    private String SystemPrompt;
    private String nextStepPrompt;

    //代理状态
    private AgentState state = AgentState.IDLE;

    //执行步骤控制
    private int currentStep = 0;
    private int maxSteps = 7;

    //LLM 模型
    private ChatClient chatClient;

    //Memory 记忆
    private List<Message> messageList = new ArrayList<>();

    /**
     * 运行代理
     * @param userPrompt 用户输入
     * @return 执行结果
     */
    public String run(String userPrompt){
        //基础校验
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }
        if (StrUtil.isBlank(userPrompt)){
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        this.state = AgentState.RUNNING;
        messageList.add(new UserMessage(userPrompt));
        List<String> results = new ArrayList<>();
        try {
            /// 循环执行
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++){
                int stepNumber = i + 1;
                log.info("Executing Step {}/{}", stepNumber, maxSteps);
                /// 单步执行
                String stepResult = step();
                String result = "Step " + stepNumber + ": " + stepResult;
                results.add(result);
            }
            /// 检查是否超出步骤限制
            if (currentStep >= maxSteps){
                this.state = AgentState.FINISHED;
                results.add("Terminated: Max steps reached. Stopping: "+ maxSteps + " steps");
            }

            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("error executing agent: ",e);
            return "执行错误" + e.getMessage();
        }
    }

    public abstract String step();

    /**
     * 运行代理(流式输出)
     * @param userPrompt 用户输入
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt){
        SseEmitter emitter = new SseEmitter(300000L);//300秒超时
        //异步执行
        CompletableFuture.runAsync(() -> {
            //基础校验
            try {
                if (this.state != AgentState.IDLE) {
                    emitter.send("result: 错误，无法从状态运行代理" + this.state);
                    emitter.complete();
                    return;
                }
                if (StrUtil.isBlank(userPrompt)){
                    emitter.send("result: 错误，不能使用空提示词运行代理");
                    emitter.complete();
                    return;
                }
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
            this.state = AgentState.RUNNING;
            messageList.add(new UserMessage(userPrompt));
            List<String> results = new ArrayList<>();
            try {
                /// 循环执行
                for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++){
                    int stepNumber = i + 1;
                    log.info("Executing Step {}/{}", stepNumber, maxSteps);
                    /// 单步执行
                    String stepResult = step();
                    String result = "Step " + stepNumber + ": " + stepResult + "\n";
                    results.add(result);
                    //发送结果
                    emitter.send(result);
                }
                /// 检查是否超出步骤限制
                if (currentStep >= maxSteps){
                    this.state = AgentState.FINISHED;
                    results.add("Terminated: Max steps reached. Stopping: "+ maxSteps + " steps");
                    emitter.send("result: Terminated: Max steps reached. Stopping: "+ maxSteps + " steps");
                }
                //正常完成
                emitter.complete();

            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("error executing agent: ",e);
                try {
                    emitter.send("result: 执行错误" + e.getMessage());
                    emitter.complete();
                } catch (IOException ex) {
                    emitter.completeWithError(ex);
                }
            }
        });

        //设置超时处理
        emitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            log.warn("SSE connect timeout. Stopping agent.");
        });
        emitter.onCompletion(() -> {
            if (state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            log.info("SSE connection completed. Stopping agent.");
        });
        return emitter;
    }
}
